import polars as pl

# 创建示例DataFrame
df1 = pl.DataFrame({
    'id': [1, 2, 3, 4],
    'name': ['Alice', 'Bob', 'Charlie', 'David'],
    'age': [25, 30, 35, 40],
    'city': ['New York', 'Los Angeles', 'Chicago', 'Houston']
})

df2 = pl.DataFrame({
    'id': [1, 2, 3, 5],
    'name': ['Alice', 'Bob', 'Charlie', 'Eve'],
    'age': [26, 31, 35, 45],
    'city': ['New York', 'San Francisco', 'Chicago', 'Austin']
})

# 设置多列主键
all_cols = ['id', 'name', 'age', 'city']
primary_keys = ['id', 'name']
value_cols = [col for col in all_cols if col not in primary_keys]  # 非主键列

# 执行全外连接，保留所有数据
joined = df1.join(df2, on=primary_keys, how='full', suffix='_right')

# 判断来源的表达式
source_expr = pl.when(pl.col(f'{primary_keys[0]}_right').is_null()).then(
    pl.lit('df1')).when(pl.col(primary_keys[0]).is_null()).then(
        pl.lit('df2')).otherwise(pl.lit('both'))

# 构建结果列表达式
result_expr = []

# 主键列保持原值（优先使用df1的值，不存在则用df2的值）
for key in primary_keys:
    result_expr.append(
        pl.coalesce(pl.col(key), pl.col(f'{key}_right')).alias(key))

# 构建结果列表达式，让 source 列排在最前面
result_expr = [source_expr.alias('source')]  # 先添加来源列

# 再添加主键列
for key in primary_keys:
    result_expr.append(
        pl.coalesce(pl.col(key), pl.col(f'{key}_right')).alias(key))

# 最后添加非主键列

# 非主键列按规则格式化：相同值显示一个，不同值显示a/b
for col in value_cols:
    result_expr.append(
        pl.when(pl.col(col) == pl.col(f'{col}_right')).then(
            pl.col(col))  # 值相同，显示原值
        .otherwise(
            pl.coalesce(pl.col(col), pl.lit('nan')).cast(str) + '/' +
            pl.coalesce(pl.col(f'{col}_right'), pl.lit('nan')).cast(str)).
        alias(col)  # 值不同或仅存在一侧，显示a/b
    )

# 选择并格式化结果
union_result = joined.select(result_expr)

print("按主键合并后的差异结果（带来源标识）：")
print(union_result)
